# -*- coding:utf-8 -*-
from PIL import Image
import torchvision.transforms as transforms
import torchvision.models as models
import torch
from torch import nn
from flask import Flask, request, jsonify

# 分类标签
cls = ['real', 'fake']
# 模型路径
model_path = './fakeFace_model.pth'


# 加载模型函数
def load_model(path):
    net = models.resnet18(pretrained=True)
    num_ftrs = net.fc.in_features
    net.fc = nn.Linear(num_ftrs, 2)  # 修改最后一层全连接层为二分类
    net.load_state_dict(torch.load(path, map_location='cpu'))
    net.eval()
    return net


# 数据转换
data_transforms = transforms.Compose([
    transforms.Resize(128),
    transforms.ToTensor(),  # 转换为张量
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])  # 标准化
])

# 创建 Flask 应用程序对象
app = Flask(__name__)

# 加载模型
device = torch.device('cpu')
model = load_model(model_path)


# 处理预测请求
@app.route('/predict', methods=['POST'])
def predict():
    if 'file' not in request.files:
        return jsonify({'code': -1, 'msg': '无文件部分'})

    file = request.files['file']
    if file.filename == '':
        return jsonify({'code': -1, 'msg': '没有选定的文件'})

    # 将上传的文件转换为 PIL 图片
    try:
        # 将上传的文件转换为 PIL 图片
        img = Image.open(file.stream)
        # 数据转换
        img = data_transforms(img)
        img = img.unsqueeze(0).to(device).float()
    except RuntimeError:
        return jsonify({'code': -1, 'msg': '无效的图像文件'})
    # 模型预测
    prob = model(img)
    predict_cls = cls[prob[0].argmax(0)]
    return jsonify({'code': 200, 'msg': predict_cls})


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=9001)